msg_tool\scripts\emote/
psb.rs

1//! Basic Handle for all emote PSB files.
2use super::rle::*;
3use crate::ext::io::*;
4use crate::ext::psb::*;
5use crate::scripts::base::*;
6use crate::types::*;
7use crate::utils::encoding::*;
8use crate::utils::files::*;
9use crate::utils::img::*;
10use anyhow::Result;
11use base64::Engine;
12use emote_psb::*;
13use libtlg_rs::*;
14use serde::{Deserialize, Serialize};
15use std::collections::HashMap;
16use std::io::{Read, Seek, Write};
17
18#[derive(Debug)]
19pub struct PsbBuilder {}
20
21impl PsbBuilder {
22    pub fn new() -> Self {
23        Self {}
24    }
25}
26
27impl ScriptBuilder for PsbBuilder {
28    fn default_encoding(&self) -> Encoding {
29        Encoding::Utf8
30    }
31
32    fn build_script(
33        &self,
34        buf: Vec<u8>,
35        _filename: &str,
36        encoding: Encoding,
37        _archive_encoding: Encoding,
38        config: &ExtraConfig,
39        _archive: Option<&Box<dyn Script>>,
40    ) -> Result<Box<dyn Script>> {
41        Ok(Box::new(Psb::new(MemReader::new(buf), encoding, config)?))
42    }
43
44    fn build_script_from_reader(
45        &self,
46        reader: Box<dyn ReadSeek>,
47        _filename: &str,
48        encoding: Encoding,
49        _archive_encoding: Encoding,
50        config: &ExtraConfig,
51        _archive: Option<&Box<dyn Script>>,
52    ) -> Result<Box<dyn Script>> {
53        Ok(Box::new(Psb::new(reader, encoding, config)?))
54    }
55
56    fn build_script_from_file(
57        &self,
58        filename: &str,
59        encoding: Encoding,
60        _archive_encoding: Encoding,
61        config: &ExtraConfig,
62        _archive: Option<&Box<dyn Script>>,
63    ) -> Result<Box<dyn Script>> {
64        let file = std::fs::File::open(filename)?;
65        let f = std::io::BufReader::new(file);
66        Ok(Box::new(Psb::new(f, encoding, config)?))
67    }
68
69    fn extensions(&self) -> &'static [&'static str] {
70        &[]
71    }
72
73    fn script_type(&self) -> &'static ScriptType {
74        &ScriptType::EmotePsb
75    }
76
77    fn is_this_format(&self, _filename: &str, buf: &[u8], buf_len: usize) -> Option<u8> {
78        if buf_len >= 4 && buf.starts_with(b"PSB\0") {
79            return Some(10);
80        }
81        None
82    }
83
84    fn can_create_file(&self) -> bool {
85        true
86    }
87
88    fn create_file<'a>(
89        &'a self,
90        filename: &'a str,
91        writer: Box<dyn WriteSeek + 'a>,
92        encoding: Encoding,
93        file_encoding: Encoding,
94        _config: &ExtraConfig,
95    ) -> Result<()> {
96        create_file(filename, writer, encoding, file_encoding)
97    }
98}
99
100#[derive(Debug)]
101pub struct Psb {
102    psb: VirtualPsbFixed,
103    encoding: Encoding,
104    config: ExtraConfig,
105}
106
107impl Psb {
108    pub fn new<R: Read + Seek>(
109        reader: R,
110        encoding: Encoding,
111        config: &ExtraConfig,
112    ) -> Result<Self> {
113        let mut psb = PsbReader::open_psb(reader)
114            .map_err(|e| anyhow::anyhow!("Failed to open psb file: {:?}", e))?;
115        let psb = psb
116            .load()
117            .map_err(|e| anyhow::anyhow!("Failed to load psb: {:?}", e))?
118            .to_psb_fixed();
119        Ok(Self {
120            psb,
121            encoding,
122            config: config.clone(),
123        })
124    }
125
126    fn output_resource(
127        &self,
128        folder_path: &std::path::PathBuf,
129        path: String,
130        data: &[u8],
131    ) -> Result<Resource> {
132        let mut res = Resource {
133            path,
134            tlg: None,
135            rle: None,
136        };
137        if self.config.psb_process_tlg && is_valid_tlg(&data) {
138            let tlg = load_tlg(MemReaderRef::new(&data))?;
139            res.tlg = Some(TlgInfo::from_tlg(&tlg, self.encoding));
140            let outtype = self.config.image_type.unwrap_or(ImageOutputType::Png);
141            res.path = {
142                let mut pb = std::path::PathBuf::from(&res.path);
143                pb.set_extension(outtype.as_ref());
144                pb.to_string_lossy().to_string()
145            };
146            let path = folder_path.join(&res.path);
147            make_sure_dir_exists(&path)?;
148            let img = ImageData {
149                width: tlg.width as u32,
150                height: tlg.height as u32,
151                color_type: match tlg.color {
152                    TlgColorType::Bgr24 => ImageColorType::Bgr,
153                    TlgColorType::Bgra32 => ImageColorType::Bgra,
154                    TlgColorType::Grayscale8 => ImageColorType::Grayscale,
155                },
156                depth: 8,
157                data: tlg.data,
158            };
159            encode_img(img, outtype, &path.to_string_lossy(), &self.config)?;
160        } else {
161            let path = folder_path.join(&res.path);
162            make_sure_dir_exists(&path)?;
163            std::fs::write(&path, data)?;
164        }
165        Ok(res)
166    }
167
168    fn output_rle_resource(
169        &self,
170        folder_path: &std::path::PathBuf,
171        path: String,
172        data: &[u8],
173        width: i64,
174        height: i64,
175    ) -> Result<Resource> {
176        let mut res = Resource {
177            path,
178            tlg: None,
179            rle: Some(RLPixelInfo { width, height }),
180        };
181        let decompressed = rl_decompress(MemReaderRef::new(data), 4, None)?;
182        let outtype = self.config.image_type.unwrap_or(ImageOutputType::Png);
183        res.path = {
184            let mut pb = std::path::PathBuf::from(&res.path);
185            pb.set_extension(outtype.as_ref());
186            pb.to_string_lossy().to_string()
187        };
188        let path = folder_path.join(&res.path);
189        make_sure_dir_exists(&path)?;
190        let img = ImageData {
191            width: width as u32,
192            height: height as u32,
193            color_type: ImageColorType::Bgra,
194            depth: 8,
195            data: decompressed,
196        };
197        encode_img(img, outtype, &path.to_string_lossy(), &self.config)?;
198        Ok(res)
199    }
200}
201
202#[derive(Debug, Deserialize, Serialize)]
203struct TlgInfo {
204    metadata: HashMap<String, String>,
205}
206
207impl TlgInfo {
208    fn from_tlg(tlg: &Tlg, encoding: Encoding) -> Self {
209        let mut metadata = HashMap::new();
210        for (k, v) in &tlg.tags {
211            let k = if let Ok(s) = decode_to_string(encoding, &k, true) {
212                s
213            } else {
214                format!(
215                    "base64:{}",
216                    base64::engine::general_purpose::STANDARD.encode(k)
217                )
218            };
219            let v = if let Ok(s) = decode_to_string(encoding, &v, true) {
220                s
221            } else {
222                format!(
223                    "base64:{}",
224                    base64::engine::general_purpose::STANDARD.encode(v)
225                )
226            };
227            metadata.insert(k, v);
228        }
229        Self { metadata }
230    }
231
232    fn to_tlg_tags(&self, encoding: Encoding) -> Result<HashMap<Vec<u8>, Vec<u8>>> {
233        let mut tags = HashMap::new();
234        for (k, v) in &self.metadata {
235            let k = if k.starts_with("base64:") {
236                base64::engine::general_purpose::STANDARD.decode(&k[7..])?
237            } else {
238                encode_string(encoding, k, false)?
239            };
240            let v = if v.starts_with("base64:") {
241                base64::engine::general_purpose::STANDARD.decode(&v[7..])?
242            } else {
243                encode_string(encoding, v, false)?
244            };
245            tags.insert(k, v);
246        }
247        Ok(tags)
248    }
249}
250
251#[derive(Debug, Deserialize, Serialize)]
252struct RLPixelInfo {
253    width: i64,
254    height: i64,
255}
256
257#[derive(Debug, Deserialize, Serialize)]
258struct Resource {
259    path: String,
260    #[serde(skip_serializing_if = "Option::is_none")]
261    tlg: Option<TlgInfo>,
262    #[serde(skip_serializing_if = "Option::is_none")]
263    rle: Option<RLPixelInfo>,
264}
265
266impl Script for Psb {
267    fn default_output_script_type(&self) -> OutputScriptType {
268        OutputScriptType::Custom
269    }
270
271    fn is_output_supported(&self, output: OutputScriptType) -> bool {
272        matches!(output, OutputScriptType::Custom)
273    }
274
275    fn default_format_type(&self) -> FormatOptions {
276        FormatOptions::None
277    }
278
279    fn custom_output_extension<'a>(&'a self) -> &'a str {
280        "json"
281    }
282
283    fn custom_export(&self, filename: &std::path::Path, encoding: Encoding) -> Result<()> {
284        let mut data = self.psb.to_json();
285        let mut resources = Vec::new();
286        let mut extra_resources = Vec::new();
287        let folder_path = {
288            let mut pb = filename.to_path_buf();
289            pb.set_extension("");
290            pb
291        };
292        for (i, data) in self.psb.resources().iter().enumerate() {
293            let i = i as u64;
294            let res_path = self.psb.root().find_resource_key(i, vec![]);
295            if let Some(path) = &res_path {
296                if path.len() >= 2 && *path.last().unwrap() == "pixel" {
297                    let pb_data = self.psb.root();
298                    let mut pb_data = &pb_data[*path.first().unwrap()];
299                    for p in path.iter().take(path.len() - 1).skip(1) {
300                        pb_data = &pb_data[*p];
301                    }
302                    let width = pb_data["width"].as_i64();
303                    let height = pb_data["height"].as_i64();
304                    let compress = pb_data["compress"].as_str();
305                    if let (Some(w), Some(h), Some(c)) = (width, height, compress) {
306                        if c == "RL" {
307                            let res_name: Vec<_> = path
308                                .iter()
309                                .take(path.len() - 1)
310                                .map(|s| s.to_string())
311                                .collect();
312                            let res_name = res_name.join("/");
313                            let res_name = sanitize_path(&res_name);
314                            let res =
315                                self.output_rle_resource(&folder_path, res_name, data, w, h)?;
316                            resources.push(res);
317                            continue;
318                        }
319                    }
320                }
321            }
322            let res_name = res_path
323                .map(|s| s.join("/"))
324                .unwrap_or(format!("res_{}", i));
325            let res_name = sanitize_path(&res_name);
326            let res = self.output_resource(&folder_path, res_name, data)?;
327            resources.push(res);
328        }
329        for (i, data) in self.psb.extra().iter().enumerate() {
330            let i = i as u64;
331            let res_name = self
332                .psb
333                .root()
334                .find_resource_key(i, vec![])
335                .map(|s| format!("extra_{}", s.join("/")))
336                .unwrap_or(format!("extra_res_{}", i));
337            let res_name = sanitize_path(&res_name);
338            let res = self.output_resource(&folder_path, res_name, data)?;
339            extra_resources.push(res);
340        }
341        data["resources"] = json::parse(&serde_json::to_string(&resources)?)?;
342        data["extra_resources"] = json::parse(&serde_json::to_string(&extra_resources)?)?;
343        let s = json::stringify_pretty(data, 2);
344        let s = encode_string(encoding, &s, false)?;
345        let mut file = std::fs::File::create(filename)?;
346        file.write_all(&s)?;
347        Ok(())
348    }
349
350    fn custom_import<'a>(
351        &'a self,
352        custom_filename: &'a str,
353        file: Box<dyn WriteSeek + 'a>,
354        encoding: Encoding,
355        output_encoding: Encoding,
356    ) -> Result<()> {
357        create_file(custom_filename, file, encoding, output_encoding)
358    }
359}
360
361fn read_resource(
362    folder_path: &std::path::PathBuf,
363    res: &Resource,
364    encoding: Encoding,
365) -> Result<Vec<u8>> {
366    if let Some(tlg) = &res.tlg {
367        let path = folder_path.join(&res.path);
368        let imgfmt = ImageOutputType::try_from(path.as_path())?;
369        let mut img = decode_img(imgfmt, &path.to_string_lossy())?;
370        if img.depth != 8 {
371            return Err(anyhow::anyhow!(
372                "Only 8-bit images are supported for TLG conversion"
373            ));
374        }
375        let color_type = match img.color_type {
376            ImageColorType::Bgr => TlgColorType::Bgr24,
377            ImageColorType::Bgra => TlgColorType::Bgra32,
378            ImageColorType::Grayscale => TlgColorType::Grayscale8,
379            ImageColorType::Rgb => {
380                convert_rgb_to_bgr(&mut img)?;
381                TlgColorType::Bgr24
382            }
383            ImageColorType::Rgba => {
384                convert_rgba_to_bgra(&mut img)?;
385                TlgColorType::Bgra32
386            }
387        };
388        let tlg = Tlg {
389            width: img.width,
390            height: img.height,
391            version: 5,
392            color: color_type,
393            data: img.data,
394            tags: tlg.to_tlg_tags(encoding)?,
395        };
396        let mut writer = MemWriter::new();
397        save_tlg(&tlg, &mut writer)?;
398        Ok(writer.into_inner())
399    } else if let Some(rle) = &res.rle {
400        let path = folder_path.join(&res.path);
401        let imgfmt = ImageOutputType::try_from(path.as_path())?;
402        let mut img = decode_img(imgfmt, &path.to_string_lossy())?;
403        if img.depth != 8 {
404            return Err(anyhow::anyhow!(
405                "Only 8-bit images are supported for RLE conversion"
406            ));
407        }
408        if img.color_type == ImageColorType::Rgba {
409            convert_rgba_to_bgra(&mut img)?;
410        } else if img.color_type == ImageColorType::Rgb {
411            convert_rgb_to_bgr(&mut img)?;
412            convert_bgr_to_bgra(&mut img)?;
413        } else if img.color_type == ImageColorType::Bgr {
414            convert_bgr_to_bgra(&mut img)?;
415        }
416        if img.color_type != ImageColorType::Bgra {
417            return Err(anyhow::anyhow!(
418                "Only BGRA images are supported for RLE conversion"
419            ));
420        }
421        if img.width as i64 != rle.width {
422            eprintln!(
423                "Warning: Image width {} does not match RLE width {}",
424                img.width, rle.width
425            );
426        }
427        if img.height as i64 != rle.height {
428            eprintln!(
429                "Warning: Image height {} does not match RLE height {}",
430                img.height, rle.height
431            );
432        }
433        let compressed = rl_compress(MemReaderRef::new(&img.data), 4)?;
434        Ok(compressed)
435    } else {
436        let path = folder_path.join(&res.path);
437        Ok(std::fs::read(&path)?)
438    }
439}
440
441fn create_file<'a>(
442    custom_filename: &'a str,
443    mut writer: Box<dyn WriteSeek + 'a>,
444    encoding: Encoding,
445    output_encoding: Encoding,
446) -> Result<()> {
447    let input = read_file(custom_filename)?;
448    let s = decode_to_string(output_encoding, &input, true)?;
449    let data = json::parse(&s)?;
450    let resources: Vec<Resource> = serde_json::from_str(&data["resources"].dump())?;
451    let extra_resources: Vec<Resource> = serde_json::from_str(&data["extra_resources"].dump())?;
452    let mut psb = VirtualPsbFixed::with_json(&data)?;
453    let folder_path = {
454        let mut pb = std::path::PathBuf::from(custom_filename);
455        pb.set_extension("");
456        pb
457    };
458    for res in resources {
459        let res = read_resource(&folder_path, &res, encoding)?;
460        psb.resources_mut().push(res);
461    }
462    for res in extra_resources {
463        let res = read_resource(&folder_path, &res, encoding)?;
464        psb.extra_mut().push(res);
465    }
466    let psb = psb.to_psb(false);
467    let psb_writer = PsbWriter::new(psb, &mut writer);
468    psb_writer
469        .finish()
470        .map_err(|e| anyhow::anyhow!("Failed to write psb: {:?}", e))?;
471    Ok(())
472}
msg_tool\scripts\emote/psb.rs

msg_tool\scripts\emote/
psb.rs
